#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  GetPromptRequestSchema,
  ListPromptsRequestSchema,
  ListResourcesRequestSchema,
  ListToolsRequestSchema,
  McpError,
  ReadResourceRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';

// 服务器信息
const SERVER_NAME = 'wjb-mcp-server-studio';
const SERVER_VERSION = '1.0.0';

// 创建服务器实例
const server = new Server(
  {
    name: SERVER_NAME,
    version: SERVER_VERSION,
  },
  {
    capabilities: {
      resources: {},
      tools: {},
      prompts: {},
    },
  }
);

// 工具列表
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'echo',
        description: 'Echo back the input text',
        inputSchema: {
          type: 'object',
          properties: {
            text: {
              type: 'string',
              description: 'Text to echo back',
            },
          },
          required: ['text'],
        },
      },
      {
        name: 'get_system_info',
        description: 'Get basic system information',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'calculate',
        description: 'Perform basic mathematical calculations',
        inputSchema: {
          type: 'object',
          properties: {
            expression: {
              type: 'string',
              description: 'Mathematical expression to evaluate (e.g., "2 + 3 * 4")',
            },
          },
          required: ['expression'],
        },
      },
    ],
  };
});

// 提示获取处理
server.setRequestHandler(GetPromptRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  switch (name) {
    case 'code-review': {
      const code = args?.['code'] as string;
      const language = (args?.['language'] as string) || 'unknown';
      
      if (!code) {
        throw new McpError(ErrorCode.InvalidParams, 'Missing required parameter: code');
      }

      return {
        description: `代码审查提示 - ${language}`,
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请对以下${language}代码进行详细审查，包括：\n1. 代码质量和最佳实践\n2. 潜在的bug和安全问题\n3. 性能优化建议\n4. 代码可读性和维护性\n\n代码：\n\`\`\`${language}\n${code}\n\`\`\``,
            },
          },
        ],
      };
    }

    case 'generate-docs': {
      const code = args?.['code'] as string;
      const format = (args?.['format'] as string) || 'markdown';
      
      if (!code) {
        throw new McpError(ErrorCode.InvalidParams, 'Missing required parameter: code');
      }

      return {
        description: `生成${format}格式文档`,
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请为以下代码生成${format}格式的详细文档，包括：\n1. 功能描述\n2. 参数说明\n3. 返回值说明\n4. 使用示例\n5. 注意事项\n\n代码：\n\`\`\`\n${code}\n\`\`\``,
            },
          },
        ],
      };
    }

    case 'explain-concept': {
      const concept = args?.['concept'] as string;
      const level = (args?.['level'] as string) || 'intermediate';
      
      if (!concept) {
        throw new McpError(ErrorCode.InvalidParams, 'Missing required parameter: concept');
      }

      const levelDescriptions = {
        beginner: '初学者友好的方式',
        intermediate: '中等难度',
        advanced: '高级深入',
      };

      return {
        description: `${levelDescriptions[level as keyof typeof levelDescriptions]}解释概念: ${concept}`,
        messages: [
          {
            role: 'user',
            content: {
              type: 'text',
              text: `请用${levelDescriptions[level as keyof typeof levelDescriptions]}解释以下概念："${concept}"\n\n请包括：\n1. 基本定义\n2. 核心原理\n3. 实际应用场景\n4. 相关概念对比\n5. 学习建议`,
            },
          },
        ],
      };
    }

    default:
      throw new McpError(ErrorCode.MethodNotFound, `Unknown prompt: ${name}`);
  }
});

// 提示列表
server.setRequestHandler(ListPromptsRequestSchema, async () => {
  return {
    prompts: [
      {
        name: 'code-review',
        description: '代码审查提示模板',
        arguments: [
          {
            name: 'code',
            description: '需要审查的代码',
            required: true,
          },
          {
            name: 'language',
            description: '编程语言',
            required: false,
          },
        ],
      },
      {
        name: 'generate-docs',
        description: '生成文档提示模板',
        arguments: [
          {
            name: 'code',
            description: '需要生成文档的代码',
            required: true,
          },
          {
            name: 'format',
            description: '文档格式 (markdown, html, plain)',
            required: false,
          },
        ],
      },
      {
        name: 'explain-concept',
        description: '概念解释提示模板',
        arguments: [
          {
            name: 'concept',
            description: '需要解释的概念',
            required: true,
          },
          {
            name: 'level',
            description: '解释级别 (beginner, intermediate, advanced)',
            required: false,
          },
        ],
      },
    ],
  };
});

// 工具调用处理
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  switch (name) {
    case 'echo': {
      const text = args?.['text'] as string;
      if (!text) {
        throw new McpError(ErrorCode.InvalidParams, 'Missing required parameter: text');
      }
      return {
        content: [
          {
            type: 'text',
            text: `Echo: ${text}`,
          },
        ],
      };
    }

    case 'get_system_info': {
      const systemInfo = {
        platform: process.platform,
        arch: process.arch,
        nodeVersion: process.version,
        uptime: process.uptime(),
        memoryUsage: process.memoryUsage(),
        timestamp: new Date().toISOString(),
      };

      return {
        content: [
          {
            type: 'text',
            text: JSON.stringify(systemInfo, null, 2),
          },
        ],
      };
    }

    case 'calculate': {
      const expression = args?.['expression'] as string;
      if (!expression) {
        throw new McpError(ErrorCode.InvalidParams, 'Missing required parameter: expression');
      }

      try {
        // 简单的数学表达式计算（仅支持基本运算符）
        const sanitizedExpression = expression.replace(/[^0-9+\-*/().\s]/g, '');
        if (sanitizedExpression !== expression) {
          throw new Error('Invalid characters in expression');
        }

        // 使用 Function 构造器安全地计算表达式
        const result = Function(`"use strict"; return (${sanitizedExpression})`)();

        return {
          content: [
            {
              type: 'text',
              text: `${expression} = ${result}`,
            },
          ],
        };
      } catch (error) {
        throw new McpError(
          ErrorCode.InternalError,
          `Failed to calculate expression: ${error instanceof Error ? error.message : 'Unknown error'}`
        );
      }
    }

    default:
      throw new McpError(ErrorCode.MethodNotFound, `Unknown tool: ${name}`);
  }
});

// 资源列表
server.setRequestHandler(ListResourcesRequestSchema, async () => {
  return {
    resources: [
      {
        uri: 'studio://server-info',
        mimeType: 'application/json',
        name: 'Server Information',
        description: 'Information about this MCP server',
      },
      {
        uri: 'studio://capabilities',
        mimeType: 'application/json',
        name: 'Server Capabilities',
        description: 'List of server capabilities and features',
      },
    ],
  };
});

// 资源读取处理
server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
  const { uri } = request.params;

  switch (uri) {
    case 'studio://server-info': {
      const serverInfo = {
        name: SERVER_NAME,
        version: SERVER_VERSION,
        description: 'A local MCP server based on Studio protocol',
        author: 'Your Name',
        capabilities: ['tools', 'resources', 'prompts'],
        uptime: process.uptime(),
        timestamp: new Date().toISOString(),
      };

      return {
        contents: [
          {
            uri,
            mimeType: 'application/json',
            text: JSON.stringify(serverInfo, null, 2),
          },
        ],
      };
    }

    case 'studio://capabilities': {
      const capabilities = {
        tools: {
          count: 3,
          available: ['echo', 'get_system_info', 'calculate'],
        },
        resources: {
          count: 2,
          available: ['studio://server-info', 'studio://capabilities'],
        },
        prompts: {
          count: 3,
          available: ['code-review', 'generate-docs', 'explain-concept'],
        },
        features: {
          stdio_transport: true,
          json_rpc: true,
          error_handling: true,
          prompts_support: true,
        },
      };

      return {
        contents: [
          {
            uri,
            mimeType: 'application/json',
            text: JSON.stringify(capabilities, null, 2),
          },
        ],
      };
    }

    default:
      throw new McpError(ErrorCode.InvalidParams, `Unknown resource: ${uri}`);
  }
});

// 启动服务器
async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);

  // 优雅关闭处理
  process.on('SIGINT', async () => {
    await server.close();
    process.exit(0);
  });

  process.on('SIGTERM', async () => {
    await server.close();
    process.exit(0);
  });
}

// 错误处理
process.on('uncaughtException', (error) => {
  console.error('Uncaught Exception:', error);
  process.exit(1);
});

process.on('unhandledRejection', (reason, promise) => {
  console.error('Unhandled Rejection at:', promise, 'reason:', reason);
  process.exit(1);
});

// 启动服务器
main().catch((error) => {
  console.error('Failed to start server:', error);
  process.exit(1);
});